import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score
from docx import Document
from docx.shared import Inches


# 一、设计指标，收集与整理数据
def load_and_preprocess_data(file_path):
    # 读取文件
    excel_file = pd.ExcelFile(file_path)

    # 获取所有表名
    sheet_names = excel_file.sheet_names
    sheet_names

    # 获取指定工作表中的数据
    df = excel_file.parse('Sheet1')

    # 查看数据的基本信息
    print('数据基本信息：')
    df.info()

    # 查看数据集行数和列数
    rows, columns = df.shape

    if rows < 100 and columns < 20:
        # 短表数据（行数少于100且列数少于20）查看全量数据信息
        print('数据全部内容信息：')
        print(df.to_csv(sep='\t', na_rep='nan'))
    else:
        # 长表数据查看数据前几行信息
        print('数据前几行内容信息：')
        print(df.head().to_csv(sep='\t', na_rep='nan'))

    return df


# 二、建立不少于两种模型来预测上市公司的营收，评价模型性能
def build_and_evaluate_models(df):
    # 提取特征和目标变量
    X = df.drop(columns=['证券代码', '证券简称', '营业收入'])
    y = df['营业收入']

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # 建立线性回归模型
    linear_model = LinearRegression()
    linear_model.fit(X_train, y_train)

    # 在线性回归模型上进行预测
    linear_y_pred = linear_model.predict(X_test)

    # 计算线性回归模型的性能指标
    linear_mse = mean_squared_error(y_test, linear_y_pred)
    linear_r2 = r2_score(y_test, linear_y_pred)

    # 建立随机森林回归模型
    rf_model = RandomForestRegressor(n_estimators=100, random_state=42)
    rf_model.fit(X_train, y_train)

    # 在随机森林回归模型上进行预测
    rf_y_pred = rf_model.predict(X_test)

    # 计算随机森林回归模型的性能指标
    rf_mse = mean_squared_error(y_test, rf_y_pred)
    rf_r2 = r2_score(y_test, rf_y_pred)

    print(f'线性回归模型均方误差: {linear_mse}, 决定系数: {linear_r2}')
    print(f'随机森林回归模型均方误差: {rf_mse}, 决定系数: {rf_r2}')

    return linear_model, rf_model, linear_mse, linear_r2, rf_mse, rf_r2


# 三、确定最终的预测模型，探索影响上市公司营收的关键因素
def select_model_and_find_features(linear_model, rf_model, linear_r2, rf_r2, df):
    if linear_r2 > rf_r2:
        final_model = linear_model
        model_name = '线性回归'
        coefficients = pd.Series(final_model.coef_, index=df.drop(
            columns=['证券代码', '证券简称', '营业收入']).columns)
        important_features = coefficients.abs().sort_values(
            ascending=False).head(5)
    else:
        final_model = rf_model
        model_name = '随机森林回归'
        importances = pd.Series(final_model.feature_importances_, index=df.drop(
            columns=['证券代码', '证券简称', '营业收入']).columns)
        important_features = importances.sort_values(ascending=False).head(5)

    print(f'最终选择的模型: {model_name}')
    print('影响上市公司营收的关键因素:')
    print(important_features)

    return final_model, model_name, important_features


# 四、生成文档撰写涵盖以上内容的实战报告
def generate_report(file_path, linear_mse, linear_r2, rf_mse, rf_r2, model_name, important_features):
    doc = Document()

    doc.add_heading('上市公司营收预测实战报告', 0)

    # 一、案例介绍
    doc.add_heading('一、案例介绍', 1)
    doc.add_paragraph('在金融领域，上市公司营收预测是非常重要的，因为它直接反映了公司的运营情况和未来发展趋势。投资者和金融数据分析师通常会关注上市公司营收预测，以评估上市公司的业绩表现和投资潜力。')

    # 二、实战目标
    doc.add_heading('二、实战目标', 1)
    doc.add_paragraph('建立上市公司营收预测模型，评价模型的性能，探索影响上市公司营收的关键因素。')

    # 三、实战过程
    doc.add_heading('三、实战过程', 1)

    # 1. 设计指标，收集与整理数据
    doc.add_heading('1. 设计指标，收集与整理数据', 2)
    doc.add_paragraph(f'从文件 {file_path} 中加载数据，并对数据进行了基本的探查。')

    # 2. 建立模型并评价性能
    doc.add_heading('2. 建立模型并评价性能', 2)
    doc.add_paragraph('建立了线性回归模型和随机森林回归模型来预测上市公司营收。')
    doc.add_paragraph(f'线性回归模型的均方误差为 {linear_mse}，决定系数为 {linear_r2}。')
    doc.add_paragraph(f'随机森林回归模型的均方误差为 {rf_mse}，决定系数为 {rf_r2}。')

    # 3. 确定最终模型和关键因素
    doc.add_heading('3. 确定最终模型和关键因素', 2)
    doc.add_paragraph(f'最终选择的模型是 {model_name} 模型。')
    doc.add_paragraph('影响上市公司营收的关键因素如下：')
    for feature, value in important_features.items():
        doc.add_paragraph(f'{feature}: {value}')

    doc.save('上市公司营收预测实战报告.docx')


if __name__ == '__main__':
    file_path = 'C:\\Users\\30299\\finance-kxd\\task3\\上市公司营收数据.xlsx'

    # 一、设计指标，收集与整理数据
    df = load_and_preprocess_data(file_path)

    # 二、建立不少于两种模型来预测上市公司的营收，评价模型性能
    linear_model, rf_model, linear_mse, linear_r2, rf_mse, rf_r2 = build_and_evaluate_models(
        df)

    # 三、确定最终的预测模型，探索影响上市公司营收的关键因素
    final_model, model_name, important_features = select_model_and_find_features(
        linear_model, rf_model, linear_r2, rf_r2, df)

    # 四、生成文档撰写涵盖以上内容的实战报告
    generate_report(file_path, linear_mse, linear_r2, rf_mse, rf_r2,
                    model_name, important_features)